package loveqq.jdbc.log;

import javassist.ClassPool;
import javassist.CtClass;
import loveqq.jdbc.log.config.AgentConfig;
import loveqq.jdbc.log.proxy.CoreProxy;

import java.lang.instrument.Instrumentation;
import java.util.HashSet;
import java.util.Properties;
import java.util.Set;

/**
 * JDBC 日志打印Log代理
 */
public class Agent {
    private static Set<String> driverProxyClass;

    static {
        try {
            // 读取driver.properties, 获取配置的Driver实现类
            Properties properties = new Properties();
            properties.load(Agent.class.getClassLoader().getResourceAsStream("driver.properties"));
            driverProxyClass = properties.stringPropertyNames();
        } catch (Exception e) {
            driverProxyClass = new HashSet<>();
        }
    }

    public static void premain(String agentArgs, Instrumentation instrumentation) throws Exception {
        // 入参配置
        AgentConfig config = AgentConfig.INSTANCE;
        config.load(agentArgs);
        if (config.isDebug()) {
            System.out.println("Agent配置: " + config);
        }

        /*
            插件思想：将JDBC整个过程的Connection，Statement全部代理，加上SQL+参数打印功能和一些辅助功能。
         */
        ClassPool pool = ClassPool.getDefault();
        CtClass preStmtClass = pool.get("java.sql.PreparedStatement");
        // 添加转换器
        instrumentation.addTransformer((loader, className, classBeingRedefined, protectionDomain, classfileBuffer) -> {
            if (className == null) {
                return classfileBuffer;
            }
            // 参数className是java/lang/String 转为 java.lang.String
            className = className.replaceAll("/", ".");
            // 需要代理的类
            CtClass ctClass = null;
            try {
                // Driver实现类代理
                if (driverProxyClass.contains(className)) {
                    ctClass = pool.get(className);
                    CoreProxy.proxyDriver(ctClass);
                    return ctClass.toBytecode();
                }
                // PreparedStatement实现类代理
                if (className.contains("PreparedStatement")) {
                    ctClass = pool.get(className);
                    if (!ctClass.isInterface() && ctClass.subtypeOf(preStmtClass)) {
                        CoreProxy.proxyStatement(ctClass);
                        return ctClass.toBytecode();
                    }
                }

            } catch (Throwable e) {
                if (config.isDebug()) {
                    e.printStackTrace();
                }
            } finally {
                if (ctClass != null) {
                    ctClass.detach();
                }
            }
            // 没有就不代理
            return classfileBuffer;
        });
        // 释放资源
        preStmtClass.detach();
    }
}
